import os

import matplotlib.pyplot as plt
import numpy as np
import torch
from torch.utils.data import DataLoader

from Dataloader import MedicalDataset
from TransUnet import TransUNet


def test(model, dataloader, device):
    model.eval()
    with torch.no_grad():
        for i, test_X in enumerate(dataloader):
            if i >= 10:  # 只显示前10张图片
                break

            # 将输入图像移到设备并进行预测
            test_X = test_X.float().to(device)
            pred_Y = model(test_X)

            # 如果模型输出是列表，选择最后一个输出
            if isinstance(pred_Y, list):
                pred_Y = pred_Y[-1]

            # 转换为NumPy格式用于显示
            test_X_np = test_X[0].permute(1, 2, 0).cpu().numpy()
            pred_Y_np = pred_Y[0].squeeze(0).cpu().numpy()
            pred_Y_np = (pred_Y_np > 0.5).astype(np.float32)  # 二值化
            # 创建画布，并添加子图
            plt.figure(figsize=(8, 4))

            # 显示输入图像
            plt.subplot(1, 2, 1)
            # 反归一化处理以显示原始图像
            mean = np.array([0.485, 0.456, 0.406])
            std = np.array([0.229, 0.224, 0.225])
            test_X_display = (test_X_np * std) + mean
            test_X_display = np.clip(test_X_display, 0, 1)
            plt.imshow(test_X_display)
            plt.title("Input Image")
            plt.xticks([])  # 去掉横坐标值
            plt.yticks([])  # 去掉纵坐标值
            plt.axis('off')

            # 显示预测的分割掩码
            plt.subplot(1, 2, 2)
            plt.imshow(pred_Y_np, cmap='gray')
            plt.title("Predicted Mask")
            plt.xticks([])  # 去掉横坐标值
            plt.yticks([])  # 去掉纵坐标值
            plt.axis('off')

            plt.tight_layout()
            plt.show()

            print(f"Processed image {i + 1}/10")


if __name__ == '__main__':

    # 创建测试数据集
    test_dataset = MedicalDataset(
        root_dir='/Volumes/For_Mac/dateset/Pulmonary_X_ray_and_masks',
        is_train=False,
        image_size=256
    )

    # 创建数据加载器
    test_loader = DataLoader(  # 注意：这里应该是test_loader而不是train_loader
        test_dataset,
        batch_size=4,
        shuffle=False,
        num_workers=4
    )

    # 初始化模型
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = TransUNet(num_classes=1).to(device)

    # 加载训练好的模型参数
    model_path = '../pt_file/TransUnet-30.pt'  # 确保这是你保存模型的正确路径
    if os.path.exists(model_path):
        model.load_state_dict(torch.load(model_path, map_location=device))
        print(f"Model loaded from {model_path}")
    else:
        print(f"Model file {model_path} does not exist.")
        exit()

    # 开始测试
    test(model, test_loader, device)
